"""
Product Quantization (PQ) 量化技术，可以使用较低的精度来表示近似向量数据，降低内存需求。对于大规模向量相似性搜索适用
"""
import os
import time

import faiss
import numpy


numpy.random.seed(0) #设置随机数种子，保证每次运行结果一致
vectors = numpy.random.rand(1000000, 768)  # 创建1000000个随机向量 , 768表示向量的维度
query = numpy.random.rand(1, 768) # 创建1个随机向量 ,用来模拟查询向量

def fun_1():
    index = faiss.IndexFlatL2(768)
    index.add(vectors)

    start = time.time()
    D, I = index.search(query, k=2)
    print(f"普通线性索引耗时：{time.time() - start}")
    print(f"最相似的向量索引：{I}")
    print(f"最相似的向量距离：{D}")

    faiss.write_index(index, "a.faiss")
    print(os.stat("a.faiss").st_size)


def fun_2():
    quantizer = faiss.IndexFlatL2(768)  # 创建一个 quantizer, 用于将向量进行量化,也就是对每个分类里的数据查询方式
    #第一个参数：量化器
    #第二个参数：向量的维度
    #第三个参数：分类数量,也就是聚类的数量
    #第四个参数：子空间的数量 , 越大越精确，但是速度越慢
    #第五个参数：子空间的分类数量，使用8位数字表示最大的值做聚类数量 ， 越大越准确，但是速度越慢
    index = faiss.IndexIVFPQ(quantizer, 768, 100, 8, 8)
    index.nprobe = 5
    index.train( vectors) #训练IVF

    index.add(vectors)

    start = time.time()
    D, I = index.search(query, k=2)
    print(f"普通线性索引耗时：{time.time() - start}")
    print(f"最相似的向量索引：{I}")
    print(f"最相似的向量距离：{D}")

    faiss.write_index(index, "b.faiss")
    print(os.stat("b.faiss").st_size)

if __name__ == '__main__':
    fun_1()
    print("---------------------------------------------")
    fun_2()